set -x

group=rtx_3
gpu=1
PORT=3323

rtx_3_p=sensecabin-4090-3
rtx_3_w=e8adf4f9-16e7-4cfc-b65e-cad12d84c108
rtx_3_r=N5IP.nn.I90.${gpu}

p="$(eval echo \${${group}_p})"
w="$(eval echo \${${group}_w})"
r="$(eval echo \${${group}_r})"
JOB_NAME=SD_INPAINTING_FOR_SAFESEAT

MODEL_NAME='/mnt/afs2d/luotianhang/cache/PretrainedModels/models--runwayml--stable-diffusion-inpainting/snapshots/afeee10def38be19995784bcc811882409d066e5'
WORK_DIR='/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/sd_inpainting/job_for_safeseat'
RESUME_MODEL='/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/sd_inpainting/job_for_safeseat/checkpoint-250000'
srun -j job_${JOB_NAME} \
     -p ${p} \
     --workspace-id ${w} \
     -N 1 \
     -f pt \
     -r ${r} \
     -o ${WORK_DIR}/nccl.log \
     -a \
     --container-mounts a79904e9-72fe-11ee-903c-4a906bc4e079:/mnt/afs2d,b9fe8ece-b9c3-11ee-a580-82381f1b95cc:/mnt/afs2d01 \
     --container-image registry.st-sh-01.sensecore.cn/cabinrd-ccr/magicube-dev:20230924-08h57m52s \
     -- bash -c "
        source /mnt/afs2d/luotianhang/start_sd.sh &&
        cp /mnt/afs2d/luotianhang/aoss.conf /root/ &&
        cp /mnt/afs2d/luotianhang/petreloss.conf /root/ &&
        cd /mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting &&
        accelerate launch  train.py \
                                    --pretrained_model_name_or_path=${MODEL_NAME} \
                                    --resolution=128  \
                                    --train_batch_size=3 --gradient_accumulation_steps=1 --gradient_checkpointing \
                                    --max_train_steps=500000 \
                                    --checkpointing_steps=500 \
                                    --checkpoints_total_limit=2 \
                                    --learning_rate=1e-05 --max_grad_norm=1 --lr_warmup_steps=0 \
                                    --mixed_precision=fp16 \
                                    --seed=423 \
                                    --num_workers=2 \
                                    --output_dir=${WORK_DIR} \
                                    --enable_xformers_memory_efficient_attention \
                                    --resume_from_checkpoint=$RESUME_MODEL \
                                    --use_8bit_adam  \
                                    --snr_gamma 5.0

     "


#Linux： 后面的值为要使用的GPU编号，正常的话是从0开始
# export CUDA_VISIBLE_DEVICES=0
# windows: 
# set CUDA_VISIBLE_DEVICES=0

# export RESUME_MODEL='/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/cabin-inpainting-model_48_tiny2_fixdata_data_4090_custom/checkpoint-9000'

# export MODEL_NAME='/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/pretrain'
# export MODEL_NAME='/mnt/afs2d/luotianhang/cache/PretrainedModels/models--runwayml--stable-diffusion-inpainting/snapshots/afeee10def38be19995784bcc811882409d066e5'
# export MODEL_NAME='/mnt/afs2d/luotianhang/cache/PretrainedModels/stable-diffusion-v1-5/stable-diffusion-v1-5'
# export MODEL_NAME='/mnt/afs2d/luotianhang/cache/PretrainedModels/models--runwayml--stable-diffusion-inpainting/snapshots/afeee10def38be19995784bcc811882409d066e5'
# export MODEL_NAME='/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/weights'
# # export MODEL_NAME='/mnt/afs2d/luotianhang/smartvehicle_diffusion/diffusers/examples/inpainting/pretrain_ppt'
# accelerate launch  train.py \
#     --pretrained_model_name_or_path=$MODEL_NAME \
#     --resolution=128  \
#     --train_batch_size=4 --gradient_accumulation_steps=1 --gradient_checkpointing \
#     --max_train_steps=250000 \
#     --checkpointing_steps=500 \
#     --checkpoints_total_limit=2 \
#     --learning_rate=1e-05 --max_grad_norm=1 --lr_warmup_steps=0 \
#     --mixed_precision=fp16 \
#     --seed=423 \
#     --num_workers=2 \
#     --output_dir=sd_inpainting \
#     --enable_xformers_memory_efficient_attention \
#     --use_8bit_adam   
#     # --resume_from_checkpoint=$RESUME_MODEL
